import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter


class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 5, padding=2)
        self.maxpool1 = nn.MaxPool2d(kernel_size=2)
        self.conv2 = nn.Conv2d(32, 32, 5, padding=2)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2)
        self.conv3 = nn.Conv2d(32, 64, 5, padding=2)
        self.maxpool3 = nn.MaxPool2d(kernel_size=2)
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(1024, 64)
        self.linear2 = nn.Linear(64, 10)

        # self.model1 = nn.Sequential(self.conv1, self.maxpool1, self.conv2, self.maxpool2, self.conv3, self.maxpool3,
        #                             self.flatten, self.linear1, self.linear2)

    def forward(self, x):
        # x = self.conv1(x)
        x = self.maxpool1(self.conv1(x))
        # x = self.conv2(x)
        x = self.maxpool2(self.conv2(x))
        # x = self.conv3(x)
        x = self.maxpool3(self.conv3(x))
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.linear2(x)
        return x


if __name__ == '__main__':
    net = SimpleNN()
    print(f"网络net：{net}\n")
    input = torch.ones((64, 3, 32, 32))
    output = net(input)
    print(f"输出output：{output}\n")
    print(f"输出output形状：{output.shape}\n")

    writer = SummaryWriter("logs/log_sifar10")
    writer.add_graph(net, input)
    writer.close()
